package com.weilus.config;

import com.nimbusds.jose.jwk.JWKSet;
import com.nimbusds.jose.jwk.RSAKey;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.security.oauth2.common.DefaultOAuth2AccessToken;
import org.springframework.security.oauth2.provider.endpoint.FrameworkEndpoint;
import org.springframework.security.oauth2.provider.token.TokenEnhancer;
import org.springframework.security.oauth2.provider.token.TokenEnhancerChain;
import org.springframework.security.oauth2.provider.token.store.JwtAccessTokenConverter;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.ResponseBody;

import java.security.*;
import java.security.interfaces.RSAPublicKey;
import java.util.Arrays;
import java.util.LinkedHashMap;
import java.util.Map;

@Configuration
public class JwtConfig {

    private KeyPair keyPair;

    public JwtConfig() throws NoSuchAlgorithmException {
        KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance("RSA");
        keyPairGenerator.initialize(2048);
        this.keyPair = keyPairGenerator.genKeyPair();
    }

    @Bean
    public JwtAccessTokenConverter accessTokenConverter(){
        JwtAccessTokenConverter converter = new JwtAccessTokenConverter();
        converter.setKeyPair(this.keyPair);
        return converter;
    }

    @Bean
    public TokenEnhancerChain enhancerChain(JwtAccessTokenConverter accessTokenConverter){
        TokenEnhancerChain enhancerChain = new TokenEnhancerChain();
        TokenEnhancer nestedTokenEnhancer = (oAuth2AccessToken, oAuth2Authentication) -> {
            DefaultOAuth2AccessToken token = new DefaultOAuth2AccessToken(oAuth2AccessToken);
            Map<String, Object> info = new LinkedHashMap(oAuth2AccessToken.getAdditionalInformation());
            info.put("iss","zhzs.lisle.cc");
            info.put("sub",oAuth2Authentication.getName());
            token.setAdditionalInformation(info);
            return token;
        };
        enhancerChain.setTokenEnhancers(Arrays.asList(nestedTokenEnhancer,accessTokenConverter));
        return enhancerChain;
    }

    public PublicKey getPubKey(){
        return this.keyPair.getPublic();
    }


    @FrameworkEndpoint
    class JwkSetUriEndpoint {
        private JwtConfig jwtConfig;
        public JwkSetUriEndpoint(JwtConfig jwtConfig){
            this.jwtConfig = jwtConfig;
        }
        @RequestMapping("/oauth/publicKey")
        public @ResponseBody Map<String,Object> publicKey(){
            RSAPublicKey rsaPublicKey = (RSAPublicKey) jwtConfig.getPubKey();
            RSAKey rsaKey = new RSAKey.Builder(rsaPublicKey).build();
            return new JWKSet(rsaKey.toPublicJWK()).toJSONObject();
        }
    }

}
